package com.swallow.auth.infrastructure.acl.oauth2.support;

import com.swallow.auth.common.utils.CustomerFuture;
import com.swallow.auth.domain.uniqno.NumWorker;
import com.swallow.auth.infrastructure.acl.oauth2.ext.ExtJdbcOAuth2AuthorizationService;
import com.swallow.auth.infrastructure.acl.oauth2.utils.SecurityUtils;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.OidcScopes;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContext;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;

import java.security.Principal;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;

/**
 * @author: yangjie.deng@resico.cn
 * @since: 2024-04-10 13:40:52
 * @version: v1.0.0
 * @describe: 处理自定义认证生成access_token、refresh_token 抽象处理类
 */
@Setter
@Getter
@Slf4j
public abstract class AbstractAuthenticationProvider<T extends Authentication> implements AuthenticationProvider {
    private OAuth2TokenGenerator<?> tokenGenerator;

    private AuthenticationManager authenticationManager;

    private OAuth2AuthorizationService authorizationService;

    private static final OAuth2TokenType ID_TOKEN_TOKEN_TYPE = new OAuth2TokenType(OidcParameterNames.ID_TOKEN);

    protected abstract AuthorizationGrantType getAuthorizationGrantType(T sourceAuthentication);

    protected OAuth2AccessTokenAuthenticationToken doBuildAccessToken(RegisteredClient registeredClient,
                                                                       Authentication authenticateRet,
                                                                       T sourceAuthentication,
                                                                       OAuth2ClientAuthenticationToken clientPrincipal) {
        // 获取请求的授权范围
        Set<String> scopes = getAuthorizedScopes(registeredClient, sourceAuthentication);
        AuthorizationServerContext context = AuthorizationServerContextHolder.getContext();

        // 以下内容摘抄自OAuth2AuthorizationCodeAuthenticationProvider
        CompletableFuture<DefaultOAuth2TokenContext.Builder> builderFuture =
                CustomerFuture.supplyAsync(() -> createTokenContextBuilder(registeredClient,
                                                                            authenticateRet,
                                                                            sourceAuthentication,
                                                                            scopes,
                                                                            context));
        // 初始化 OAuth2Authorization
        CompletableFuture<OAuth2Authorization.Builder> authorizationBuilderFuture =
                CustomerFuture.supplyAsync(() -> createAuthorizationBuilder(registeredClient,
                                                                             authenticateRet,
                                                                             sourceAuthentication,
                                                                             scopes));

        // ----- Access token -----
        DefaultOAuth2TokenContext.Builder tokenContextBuilder = builderFuture.join();
        OAuth2TokenContext tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.ACCESS_TOKEN).build();
        OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
        if (generatedAccessToken == null) {
            OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,"The token generator failed to generate the access token.", "");
            throw new OAuth2AuthenticationException(error);
        }
        OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
                                                              generatedAccessToken.getTokenValue(),
                                                              generatedAccessToken.getIssuedAt(),
                                                              generatedAccessToken.getExpiresAt(),
                                                              tokenContext.getAuthorizedScopes());

        OAuth2Authorization.Builder authorizationBuilder = authorizationBuilderFuture.join();
        if (generatedAccessToken instanceof ClaimAccessor claimAccessor) {
            authorizationBuilder.token(accessToken, metadata ->
                    metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, claimAccessor.getClaims()));
        } else {
            authorizationBuilder.accessToken(accessToken);
        }

        // ----- Refresh token -----
        OAuth2RefreshToken refreshToken = null;
        if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN) &&
                // Do not issue refresh token to public client
                !clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {

            tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build();
            OAuth2Token generatedRefreshToken = this.tokenGenerator.generate(tokenContext);

            if (generatedRefreshToken instanceof OAuth2RefreshToken token) {
                refreshToken = token;
                authorizationBuilder.refreshToken(token);
            } else {
                SecurityUtils.throwError(OAuth2ErrorCodes.SERVER_ERROR, "The token generator failed to generate the refresh token.");
            }
        }

        // 参数传递
        Map<String, Object> additionalParameters = new HashMap<>(1);
        // ----- ID token -----
        if (scopes.contains(OidcScopes.OPENID)) {
            // ID token customizer may need access to the access token and/or refresh token
            tokenContext = tokenContextBuilder.tokenType(ID_TOKEN_TOKEN_TYPE).authorization(authorizationBuilder.build()).build();
            // @formatter:on
            OAuth2Token generatedIdToken = this.tokenGenerator.generate(tokenContext);
            if (generatedIdToken instanceof Jwt jwt) {
                OidcIdToken idToken = new OidcIdToken(generatedIdToken.getTokenValue(), generatedIdToken.getIssuedAt(),
                        generatedIdToken.getExpiresAt(), jwt.getClaims());
                authorizationBuilder.token(idToken, (metadata) -> metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, idToken.getClaims()));

                // 放入idToken
                additionalParameters.put(OidcParameterNames.ID_TOKEN, idToken.getTokenValue());
            } else {
                SecurityUtils.throwError(OAuth2ErrorCodes.SERVER_ERROR, "The token generator failed to generate the ID token.");
            }
        }
        // 生成实例
        OAuth2Authorization authorization = authorizationBuilder.build();
        // 保存
        this.authorizationService.save(authorization);

        if (authorizationService instanceof ExtJdbcOAuth2AuthorizationService service) {
            // 挤下线功能, 需要则可注释掉该代码
            CustomerFuture.runAsync(() -> service.deleteByAuthInfo(authorization));
        }

        return new OAuth2AccessTokenAuthenticationToken(registeredClient, clientPrincipal, accessToken, refreshToken, additionalParameters);
        
    }

    private OAuth2Authorization.Builder createAuthorizationBuilder(RegisteredClient registeredClient,
                                                                   Authentication authenticateRet,
                                                                   T sourceAuthentication,
                                                                   Set<String> scopes) {
        return OAuth2Authorization.withRegisteredClient(registeredClient)
                                  // 生成ID,指定ID生成规则,不使用默认的UUID
                                  .id(NumWorker.bizNo())
                                  // 加入当前用户认证信息，防止刷新token时因获取不到认证信息而抛出空指针异常
                                  // 存入授权scope
                                  .authorizedScopes(scopes)
                                  // 当前授权用户名称
                                  .principalName(authenticateRet.getName())
                                  // 设置当前用户认证信息
                                  .attribute(Principal.class.getName(), authenticateRet)
                                  .authorizationGrantType(getAuthorizationGrantType(sourceAuthentication));
    }

    private DefaultOAuth2TokenContext.Builder createTokenContextBuilder(RegisteredClient registeredClient,
                                                                        Authentication authenticateRet,
                                                                        T sourceAuthentication,
                                                                        Set<String> scopes,
                                                                        AuthorizationServerContext authorizationServerContext) {
        return DefaultOAuth2TokenContext.builder()
                                        .registeredClient(registeredClient)
                                        .principal(authenticateRet)
                                        .authorizationServerContext(authorizationServerContext)
                                        .authorizedScopes(scopes)
                                        .authorizationGrantType( getAuthorizationGrantType(sourceAuthentication) )
                                        .authorizationGrant(sourceAuthentication);
    }

    protected abstract Set<String> getAuthorizedScopes(RegisteredClient registeredClient, T sourceAuthentication);


}
